package it.uniroma2.drp;

import it.uniroma2.dtk.dt.GenericDT;

import it.uniroma2.svd.ExecSVD;
import it.uniroma2.svd.writer.DenseBinaryMatrix;
import it.uniroma2.util.math.ArrayMath;
import it.uniroma2.util.math.MatrixUtils;
import it.uniroma2.util.tree.Tree;


import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Arrays;
import java.util.Comparator;
import java.util.Date;
import java.util.List;
import java.util.Vector;

public class Learner {

	private GenericDT ts = null;
	private InputVectorGenerator ivg = null;
	private int type = 0;
	private File drp_parser_output_file = null;
	private File drp_parser_input_file = null;
	
	private File learnt_drp = null;
	
	
	/**
	 * 
	 * @param argv
	 *     argv[0] : file or directory of files with trees
	 *     argv[1] : ouput_dir
	 *     argv[2] : type ( 0 (unigrams), 1 (unigrams+bigrams), 2 (unigrams+bigrams+trigrams), ...)
	 *     argv[3] : dimension of the space
	 *     argv[4] : max number of eigenvalues (k) for SVD
	 *     argv[5] : configuration file
	 * @throws Exception
	 */
	public static void main(String [] argv) throws Exception {
		try {
			File treeFile = new File(argv[0]);
			File output_dir = new File(argv[1]);
			int type = Integer.parseInt(argv[2]);
			int space_dim =  Integer.parseInt(argv[3]);
			int k =  Integer.parseInt(argv[4]);
			File configuration_file = new File(argv[5]); 
			
			GenericDT fft = new GenericDT(0, space_dim, false, false, GenericDT.Types.CONVOLUTION);
			fft.setLambda(0);
			Learner lmb = new Learner(fft,type);

			// First step: Preparing the matrices for Singular Value Decomposition

			lmb.processTreeFile(treeFile,output_dir);

			
			// Second step: Executing the external (c) svd
			
			File svd_base_output_dir = new File(output_dir,treeFile.getName() + "_" + lmb.type);
			int existing_lower_k = lmb.selectLowerK(svd_base_output_dir, k);
		
			if (existing_lower_k == -1) {
				ExecSVD svd = new ExecSVD(k, svd_base_output_dir, 
						lmb.drp_parser_input_file,
						true, true);
				svd.runSVD();
				existing_lower_k = k;
			}
						
			// Final matrix multiplication
			
			lmb.finalParserMatrixGeneration(svd_base_output_dir,k,existing_lower_k);
			System.out.println("DRP learnt!");
			

			lmb.writeConfigFile(configuration_file);
			
			
			
		} catch (ArrayIndexOutOfBoundsException e) {
			e.printStackTrace();
			System.out.println("USE:\n" +
			 " argv[0] : file or directory of files with trees\n"+
			 " argv[1] : ouput_dir\n"+
			 " argv[2] : type ( 0 (unigrams), 1 (unigrams+bigrams), 2 (unigrams+bigrams+trigrams), ...)\n"+
			 " argv[3] : dimension of the space\n"+
			 " argv[4] : max number of eigenvalues (k) for SVD\n" +
			 " argv[5] : configuration file\n");
		}
	}
	

	
	private void writeConfigFile(File configuration_file) throws IOException {
		System.out.println("Actual configuration in:\n\t" + configuration_file.getAbsolutePath());
		BufferedWriter conf = new BufferedWriter(new FileWriter(configuration_file));
		conf.write("Distributed Representation Parser\nConfiguration File\nRun: "+ (new Date(System.currentTimeMillis())).toString() +"\n");
		conf.write("DRP\t"+ learnt_drp.getAbsolutePath() +"\n");
		conf.close();
		
	}

	private int selectLowerK(File svd_repository,int expected_k) {
		int k_min_above_expected_k = -1;
		if (svd_repository.exists() && svd_repository.isDirectory()) {
			File [] k_list = svd_repository.listFiles();
			for (File kf:k_list) {
				try {
					int k = new Integer(kf.getName());
					if ((k >= expected_k) && (k < k_min_above_expected_k || k_min_above_expected_k==-1) ) k_min_above_expected_k = k;
				} catch (java.lang.NumberFormatException e) {
					
				}
			}
		}
		return k_min_above_expected_k;
	}

	public void finalParserMatrixGeneration(File svd_file_base,int k, int existing_lower_k) throws Exception {
		
		File svd_files = new File(svd_file_base, "" +existing_lower_k); 

		DenseBinaryMatrix<Float> input_dt = new DenseBinaryMatrix<Float>(Float.class);
		input_dt.openFile(drp_parser_input_file.getAbsolutePath(), "r");
		input_dt.readSize();

		System.out.println("DT  " + input_dt.getRows() + "x" + input_dt.getCols());
		Float[][] DT = input_dt.getFullMatrixFloat();
		


		BufferedReader sigma = new BufferedReader(new FileReader(new File(svd_files,"svd-S")));
		String eigv = sigma.readLine();
		int found_k_values = new Integer(eigv);
		eigv = sigma.readLine();
		
		Float [] sigma_inv = new Float[found_k_values];
		int j = 0;
		while (eigv!=null) {
			sigma_inv[j] = 1/(new Float(eigv));
			j++;
			eigv = sigma.readLine();
		}

		
		DenseBinaryMatrix<Float> output_Ut = new DenseBinaryMatrix<Float>(Float.class);
		output_Ut.openFile((new File(svd_files,"svd-Ut")).getAbsolutePath(), "r");
		output_Ut.readSize();
		System.out.println("Ut  " + output_Ut.getRows() + "x" + output_Ut.getCols());
		System.out.println("required k     = " + k);
		System.out.println("k values found = " + found_k_values);

		Float[][] SigmaInvUtDT = new Float[found_k_values][input_dt.getCols()];

		output_Ut.seekRowZero();
		for (int i=0;i < found_k_values;i++) {
			
			System.out.print("."+i);
			Float[] rowUt_ith = output_Ut.getNextFullRowFloat();
			ArrayMath.scalardot(sigma_inv[i], rowUt_ith);
			SigmaInvUtDT[i] = MatrixUtils.multiply(rowUt_ith,DT);
		}
		
		System.out.println("SigmaInvUtDT " + SigmaInvUtDT.length + " " + SigmaInvUtDT[0].length);

		
		DenseBinaryMatrix<Float> output_Vt = new DenseBinaryMatrix<Float>(Float.class);
		output_Vt.openFile((new File(svd_files,"svd-Vt")).getAbsolutePath(), "r");
		output_Vt.readSize();
		
		Float[][] Vt = new Float[found_k_values][output_Vt.getCols()];
		for (int i=0;i < found_k_values;i++) Vt[i] = output_Vt.getNextFullRowFloat();
		
		System.out.println("Original Vt " + output_Vt.getRows() + "x" + output_Vt.getCols());
		System.out.println("Used Vt     " + Vt.length + "x" + Vt[0].length );
		
		Float[][] VSigmaInvUtDT = MatrixUtils.multiply(Vt, SigmaInvUtDT, true, false);
		
		System.out.println("Writing final matrix");
		learnt_drp = (new File(svd_file_base,"drp"+  found_k_values + ".dbm"));
		DenseBinaryMatrix<Float> drp = new DenseBinaryMatrix<Float>(Float.class);
		drp.openFile(learnt_drp.getAbsolutePath(), "rw");
		drp.setFullMatrixFloat(VSigmaInvUtDT);
		drp.closeFile();
	
	}


	private void processTreeFile(File treeFile, File ouput_dir) throws Exception {
		DenseBinaryMatrix<Float> output_dt = null;
		DenseBinaryMatrix<Float> output_input = null;
		
		drp_parser_input_file = new File(ouput_dir,treeFile.getName() + "_" + type +".dbm");
		
		if (!(drp_parser_input_file.exists())) {
			output_input = new DenseBinaryMatrix<Float>(Float.class);
			output_input.openFile(drp_parser_input_file.getAbsolutePath(), "rw");
			output_input.setCol(ts.getVectorSize());
		} else {
			System.out.println("--- " + type +"-POS-gram matrix already exists --- \n\t"+ drp_parser_input_file.getAbsolutePath());
		}
		
		drp_parser_output_file = new File(ouput_dir,treeFile.getName() + "_dt.dbm");

		
		if (!drp_parser_output_file.exists()) {
			output_dt = new DenseBinaryMatrix<Float>(Float.class);
			output_dt.openFile(drp_parser_output_file.getAbsolutePath(),"rw");
			output_dt.setCol(ts.getVectorSize());
		} else {
			System.out.println("--- Distributed Tree Matrix already exists ---\n\t"+
					drp_parser_output_file.getAbsolutePath());
		}
		
		if (output_input!=null) {
			int rows = processTreeFile(treeFile,output_input,output_dt);
			if (output_dt!=null) {
				output_dt.setRows(rows);
				output_dt.writeSize();
				output_dt.closeFile();
				
			}
			output_input.setRows(rows);
			output_input.writeSize();
			output_input.closeFile();
		}
		
	}


  public static File[] dirAlphaOrder(File folder) {
	    if (!folder.isDirectory()) {
	      return null;
	    }
	    File files[] = folder.listFiles();
	    Arrays.sort( files, new Comparator<File>()
	    {
	      public int compare(final File o1, final File o2) {
	        return o1.getName().compareTo(o2.getName());
	      }
	    }); 
	    return files;
	  }  
	
	private int processTreeFile(File treeFile, 
			DenseBinaryMatrix<Float> output_input,
			DenseBinaryMatrix<Float> output_dt) throws Exception {
		int rows = 0;
		if (treeFile.isDirectory()) {
			System.out.println("Processing Dir  > " + treeFile.getName());
			System.out.println(" About to process these files: " + treeFile.getName());
			for (File f:dirAlphaOrder(treeFile))	System.out.print(" " + f.getName());
			System.out.println();
			for (File f:dirAlphaOrder(treeFile))	{
				rows += processTreeFile(f,output_input,output_dt);
			}
			System.out.println();
		} else {
			System.out.print("." + treeFile.getName());
			BufferedReader in = new BufferedReader(new FileReader(treeFile));
			String parsed_sentence = loadPennTree(in);
			int row = 0;
			while (parsed_sentence != null) {
				//System.out.print(".");
				//System.out.println("> " + parsed_sentence);
				Tree t = Tree.fromPennTree(parsed_sentence);
				t = polishPennTree(t);

				if (output_dt != null) {
					output_dt.setFullRowFloat(row, convert(ts.dt(t)));
					//System.out.print("+");
				}
				output_input.setFullRowFloat(row, convert(ivg.generate(t)));
				row++;
				parsed_sentence = loadPennTree(in);
			}
			rows = row;
			in.close();
		}
		return rows;
	}

	public Learner(GenericDT dtfGenerator,int type) {
		ts = dtfGenerator;
		this.type = type;
		ivg = new InputVectorGenerator(type,dtfGenerator);
	}
	
	
	private Float [] convert(double [] v) {
		Float [] f = new Float[v.length];
		for (int i=0;i < v.length;i++) f[i] = new Float(v[i]);
		return f;
	}
	
	public static  String loadPennTree(BufferedReader in) throws IOException{
		String input = in.readLine();
		String output = null;
		int open_pars = 0;
		while (input != null && input.trim().equals("")) input = in.readLine();
		if (input!=null)  {
			open_pars = count_chars(input,'(') - count_chars(input,')');
			output = input.trim();
		}
		while (input != null && open_pars > 0) {
			input = in.readLine();
			open_pars += count_chars(input,'(') - count_chars(input,')');
			output += " " + input.trim();
		}
		return output;
	}


	private Tree polishPennTree(Tree c) {
		Tree t  = null;
		if (isPreterminal(c)) {
			
			//System.out.println(" +++ " + c.getType());
			if (!(c.getRootLabel().trim().contains("*"))) t = c; 
//			else {
//				System.out.println(" >>> " + c.getType());
//			}
		} else {
			List<Tree> cc  = c.getChildren();

			Vector<Tree> new_cc = new Vector<Tree>();
			for (Tree o:cc) {
				Tree new_o = polishPennTree(o);
				if (new_o!=null) new_cc.add(o);
			}
			if (new_cc.size() != 0) {
				c.setChildren(new_cc); 
				t = c;
			}
		}
		if (t != null) {
			if (t.getRootLabel().contains("-")) c.setRootLabel(c.getRootLabel().substring(0,c.getRootLabel().indexOf('-')) ); 
		}
		return t;
	}

	public boolean isPreterminal(Tree t) {
		return t.getChildren().size() == 1 && t.getChildren().get(0).isTerminal();
	}			

	private static  int count_chars(String input, char c) {
		int count = 0;
		
		for (int i = 0;i < input.length(); i++) {
			count += (input.charAt(i)==c? 1 : 0);
		}
		
		return count;
	}


	
}
